import sun.misc.Unsafe;

import java.util.concurrent.CountDownLatch;
import java.util.stream.IntStream;

/**
 * 自定义 锁
 */
public class LockDemo {
    private volatile int state;
    // Unsafe实例
    private static final  Unsafe unsafe;
    // value字段的偏移量
    private static final long stateOffset;
    // value字段的偏移量
    private static final long tailOffset;

    static {
        try {
            unsafe = sun.misc.Unsafe.getUnsafe();
            Class<?> ak = LockDemo.class;
            stateOffset = unsafe.objectFieldOffset(ak.getDeclaredField("state"));
            Class<?> bk = LockDemo.Node.class;
            tailOffset = unsafe.objectFieldOffset(bk.getDeclaredField("tail"));
        } catch (Exception e) {
            throw new Error(e);
        }
    }

    private boolean compareAndSetState(int expect, int update) {
        return unsafe.compareAndSwapInt(this, stateOffset, expect, update);
    }

    private static class Node{
        /** 存储的元素为线程 */
        Thread thread;
        /** 前一个节点（可以没有，但实现起来很困难） */
        Node prev;
        /** 后一个节点 */
        Node next;

        public Node(){}

        public Node(Thread thread, Node prev){
            this.thread = thread;
            this.prev = prev;
        }
    }

    /** 链表头 */
    private volatile Node head;
    /** 链表尾 */
    private volatile Node tail;

    /**
     * 原子更新链表tail字段
     * @return
     */
    private boolean compareAndSetTail(Node expect, Node update){
        return unsafe.compareAndSwapObject(this, tailOffset, expect, update);
    }


    /**
     * 加锁
     * 1. 尝试获取锁，成功了就直接返回；
     * 2. 未获取到锁，就进入队列排队；
     * 3. 入队之后，再次尝试获取锁；
     * 4. 如果不成功，就阻塞；
     * 5. 如果成功了，就把头节点后移一位，并清空当前节点的内容，且与上一个节点断绝关系；
     * 6. 加锁结束；
     */
    public void lock(){
        // 尝试更新state字段，更新成功说明占有了锁
        if (compareAndSetState(0, 1)) {
            return;
        }

        // 未更新成功则入队
        Node node = enqueue();
        Node prev = node.prev;

        // 再次尝试获取锁，需要检测上一个节点是不是head，按入队顺序加锁
        while (node.prev != head || !compareAndSetState(0, 1)) {
            // 未获取到锁，阻塞
            unsafe.park(false, 0L);
        }
        // 下面不需要原子更新，因为同时只有一个线程访问到这里
        // 获取到锁了且上一个节点是head
        // head后移一位
        head = node;
        // 清空当前节点的内容，协助GC
        node.thread = null;
        // 将上一个节点从链表中剔除，协助GC
        node.prev = null;
        prev.next = null;
    }

    /**
     * 入队
     * @return
     */
    private Node enqueue(){
        while(true){
            // 获取尾节点
            Node t = tail;
            // 构造新节点
            Node node = new Node(Thread.currentThread(), t);
            // 不断尝试原子更新尾节点
            if (compareAndSetTail(t, node)) {
                // 更新尾节点成功了，让原尾节点的next指针指向当前节点
                t.next = node;
                return node;
            }
        }
    }

    /**
     * 解锁
     * 1. 把state改成0，这里不需要CAS更新，因为现在还在加锁中，只有一个线程去更新，在这句之后就释放了锁；
     * 2. 如果有下一个节点就唤醒它；
     * 3. 唤醒之后就会接着走上面lock()方法的while循环再去尝试获取锁；
     * 4. 唤醒的线程不是百分之百能获取到锁的，因为这里state更新成0的时候就解锁了，之后可能就有线程去尝试加锁了。
     */
    public void unlock(){
        // 把state更新成0，这里不需要原子更新，因为同时只有一个线程访问到这里
        state = 0;
        // 下一个待唤醒的节点
        Node next = head.next;
        // 下一个节点不为空，就唤醒它
        if (next != null) {
            unsafe.unpark(next.thread);
        }
    }


    private static int count = 0;

    public static void main(String[] args) throws InterruptedException {
        LockDemo lock = new LockDemo();

        CountDownLatch countDownLatch = new CountDownLatch(1000);

        IntStream.range(0, 1000).forEach(i -> new Thread(() -> {
            lock.lock();

            try {
                IntStream.range(0, 10000).forEach(j -> {
                    count++;
                });
            } finally {
                lock.unlock();
            }
            countDownLatch.countDown();
        }, "tt-" + i).start());

        countDownLatch.await();

        System.out.println(count);
    }
}
